#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
图像预处理效果对比工具
生成多种预处理效果供肉眼观察和对比
"""

import cv2
import numpy as np
import os
import sys
from datetime import datetime

class PreprocessingComparison:
    """图像预处理效果对比工具"""
    
    def __init__(self):
        self.methods = {
            'original': self.original,
            'grayscale': self.grayscale,
            'contrast_enhanced': self.contrast_enhanced,
            'histogram_equalized': self.histogram_equalized,
            'clahe': self.clahe,
            'gaussian_blur': self.gaussian_blur,
            'bilateral_filter': self.bilateral_filter,
            'adaptive_threshold': self.adaptive_threshold,
            'otsu_threshold': self.otsu_threshold,
            'morphology_opening': self.morphology_opening,
            'morphology_closing': self.morphology_closing,
            'edge_enhancement': self.edge_enhancement,
            'unsharp_mask': self.unsharp_mask,
            'gamma_correction_bright': self.gamma_correction_bright,
            'gamma_correction_dark': self.gamma_correction_dark,
            'denoising': self.denoising,
            'combined_enhancement': self.combined_enhancement,
            'text_optimized': self.text_optimized
        }
    
    def original(self, image):
        """原始图像"""
        return image
    
    def grayscale(self, image):
        """灰度转换"""
        if len(image.shape) == 3:
            return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        return image
    
    def contrast_enhanced(self, image):
        """对比度增强"""
        gray = self.grayscale(image)
        return cv2.convertScaleAbs(gray, alpha=2.0, beta=0)
    
    def histogram_equalized(self, image):
        """直方图均衡化"""
        gray = self.grayscale(image)
        return cv2.equalizeHist(gray)
    
    def clahe(self, image):
        """CLAHE对比度限制自适应直方图均衡化"""
        gray = self.grayscale(image)
        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
        return clahe.apply(gray)
    
    def gaussian_blur(self, image):
        """高斯模糊降噪"""
        gray = self.grayscale(image)
        return cv2.GaussianBlur(gray, (5, 5), 0)
    
    def bilateral_filter(self, image):
        """双边滤波（保边降噪）"""
        gray = self.grayscale(image)
        return cv2.bilateralFilter(gray, 9, 75, 75)
    
    def adaptive_threshold(self, image):
        """自适应阈值二值化"""
        gray = self.grayscale(image)
        return cv2.adaptiveThreshold(
            gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
            cv2.THRESH_BINARY, 11, 2
        )
    
    def otsu_threshold(self, image):
        """Otsu阈值二值化"""
        gray = self.grayscale(image)
        _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        return binary
    
    def morphology_opening(self, image):
        """形态学开运算"""
        gray = self.grayscale(image)
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
        return cv2.morphologyEx(gray, cv2.MORPH_OPEN, kernel)
    
    def morphology_closing(self, image):
        """形态学闭运算"""
        gray = self.grayscale(image)
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
        return cv2.morphologyEx(gray, cv2.MORPH_CLOSE, kernel)
    
    def edge_enhancement(self, image):
        """边缘增强"""
        gray = self.grayscale(image)
        # 使用拉普拉斯算子
        laplacian = cv2.Laplacian(gray, cv2.CV_64F)
        enhanced = gray + 0.3 * laplacian
        return np.clip(enhanced, 0, 255).astype(np.uint8)
    
    def unsharp_mask(self, image):
        """锐化滤波"""
        gray = self.grayscale(image)
        gaussian = cv2.GaussianBlur(gray, (9, 9), 10.0)
        unsharp = cv2.addWeighted(gray, 1.5, gaussian, -0.5, 0)
        return unsharp
    
    def gamma_correction_bright(self, image):
        """伽马校正（提亮）"""
        gray = self.grayscale(image)
        gamma = 0.7  # < 1 提亮
        table = np.array([((i / 255.0) ** gamma) * 255 for i in np.arange(0, 256)]).astype("uint8")
        return cv2.LUT(gray, table)
    
    def gamma_correction_dark(self, image):
        """伽马校正（变暗）"""
        gray = self.grayscale(image)
        gamma = 1.5  # > 1 变暗
        table = np.array([((i / 255.0) ** gamma) * 255 for i in np.arange(0, 256)]).astype("uint8")
        return cv2.LUT(gray, table)
    
    def denoising(self, image):
        """去噪"""
        gray = self.grayscale(image)
        return cv2.fastNlMeansDenoising(gray)
    
    def combined_enhancement(self, image):
        """组合增强（CLAHE + 双边滤波 + 轻微锐化）"""
        gray = self.grayscale(image)
        
        # CLAHE增强对比度
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        enhanced = clahe.apply(gray)
        
        # 双边滤波降噪
        filtered = cv2.bilateralFilter(enhanced, 5, 50, 50)
        
        # 轻微锐化
        gaussian = cv2.GaussianBlur(filtered, (3, 3), 1.0)
        sharpened = cv2.addWeighted(filtered, 1.2, gaussian, -0.2, 0)
        
        return sharpened
    
    def text_optimized(self, image):
        """文本识别优化（专门针对OCR）"""
        gray = self.grayscale(image)
        
        # 1. 适度对比度增强
        enhanced = cv2.convertScaleAbs(gray, alpha=1.2, beta=10)
        
        # 2. 降噪但保持边缘
        denoised = cv2.bilateralFilter(enhanced, 5, 80, 80)
        
        # 3. 自适应阈值
        binary = cv2.adaptiveThreshold(
            denoised, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
            cv2.THRESH_BINARY, 9, 2
        )
        
        # 4. 轻微形态学操作连接文字
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (2, 1))
        processed = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
        
        return processed
    
    def process_image(self, image_path, output_dir="preprocessing_results"):
        """处理图像并生成所有预处理版本"""
        
        if not os.path.exists(image_path):
            print(f"❌ 图像文件不存在: {image_path}")
            return
        
        # 创建输出目录
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        # 读取原始图像
        original_image = cv2.imread(image_path)
        if original_image is None:
            print(f"❌ 无法读取图像: {image_path}")
            return
        
        print(f"📷 正在处理图像: {os.path.basename(image_path)}")
        print(f"📐 原始尺寸: {original_image.shape[1]}x{original_image.shape[0]}")
        print(f"📂 输出目录: {output_dir}")
        print()
        
        # 生成所有预处理版本
        results = {}
        for method_name, method_func in self.methods.items():
            try:
                print(f"🔄 处理中: {method_name}")
                
                processed = method_func(original_image)
                
                # 保存处理后的图像
                filename = f"{method_name}_{os.path.splitext(os.path.basename(image_path))[0]}.jpg"
                output_path = os.path.join(output_dir, filename)
                
                # 确保是单通道图像时转换为三通道保存
                if len(processed.shape) == 2:
                    processed_color = cv2.cvtColor(processed, cv2.COLOR_GRAY2BGR)
                else:
                    processed_color = processed
                
                cv2.imwrite(output_path, processed_color)
                results[method_name] = output_path
                
                print(f"   ✅ 保存至: {filename}")
                
            except Exception as e:
                print(f"   ❌ 处理失败: {e}")
        
        print(f"\n🎯 预处理完成！")
        print(f"📊 成功生成 {len(results)} 个预处理版本")
        print(f"📁 所有文件保存在: {output_dir}")
        
        # 生成对比HTML文件
        self.generate_comparison_html(results, output_dir, os.path.basename(image_path))
        
        return results
    
    def generate_comparison_html(self, results, output_dir, original_filename):
        """生成HTML对比页面"""
        html_content = f"""
<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>图像预处理效果对比 - {original_filename}</title>
    <style>
        body {{
            font-family: Arial, sans-serif;
            margin: 20px;
            background-color: #f5f5f5;
        }}
        .header {{
            text-align: center;
            margin-bottom: 30px;
        }}
        .grid {{
            display: grid;
            grid-template-columns: repeat(auto-fit, minmax(300px, 1fr));
            gap: 20px;
        }}
        .image-item {{
            background: white;
            border-radius: 8px;
            padding: 15px;
            box-shadow: 0 2px 8px rgba(0,0,0,0.1);
            text-align: center;
        }}
        .image-item h3 {{
            margin: 0 0 10px 0;
            color: #333;
            font-size: 14px;
        }}
        .image-item img {{
            max-width: 100%;
            height: auto;
            border: 1px solid #ddd;
            border-radius: 4px;
        }}
        .original {{
            border: 3px solid #007bff;
        }}
        .method-description {{
            font-size: 12px;
            color: #666;
            margin-top: 8px;
            line-height: 1.4;
        }}
    </style>
</head>
<body>
    <div class="header">
        <h1>图像预处理效果对比</h1>
        <p>原始图像: {original_filename}</p>
        <p>生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}</p>
    </div>
    <div class="grid">
"""
        
        # 方法描述
        descriptions = {
            'original': '原始图像，未经任何处理',
            'grayscale': '转换为灰度图像',
            'contrast_enhanced': '对比度增强 (alpha=2.0)',
            'histogram_equalized': '直方图均衡化',
            'clahe': 'CLAHE对比度限制自适应直方图均衡化',
            'gaussian_blur': '高斯模糊降噪',
            'bilateral_filter': '双边滤波保边降噪',
            'adaptive_threshold': '自适应阈值二值化',
            'otsu_threshold': 'Otsu自动阈值二值化',
            'morphology_opening': '形态学开运算',
            'morphology_closing': '形态学闭运算',
            'edge_enhancement': '拉普拉斯边缘增强',
            'unsharp_mask': '锐化滤波',
            'gamma_correction_bright': '伽马校正提亮 (γ=0.7)',
            'gamma_correction_dark': '伽马校正变暗 (γ=1.5)',
            'denoising': '非局部均值去噪',
            'combined_enhancement': 'CLAHE + 双边滤波 + 锐化组合',
            'text_optimized': 'OCR文本识别专用优化'
        }
        
        # 确保原始图像排在第一个
        ordered_methods = ['original'] + [k for k in results.keys() if k != 'original']
        
        for method in ordered_methods:
            if method in results:
                css_class = 'original' if method == 'original' else ''
                description = descriptions.get(method, '无描述')
                filename = os.path.basename(results[method])
                
                html_content += f"""
        <div class="image-item {css_class}">
            <h3>{method.replace('_', ' ').title()}</h3>
            <img src="{filename}" alt="{method}" onclick="window.open('{filename}', '_blank')">
            <div class="method-description">{description}</div>
        </div>
"""
        
        html_content += """
    </div>
    <div style="text-align: center; margin-top: 30px; color: #666;">
        <p>点击图像可查看原始大小</p>
        <p>建议选择文字清晰、对比度好、噪声少的预处理方法</p>
    </div>
</body>
</html>
"""
        
        html_file = os.path.join(output_dir, "comparison.html")
        with open(html_file, 'w', encoding='utf-8') as f:
            f.write(html_content)
        
        print(f"🌐 生成对比页面: comparison.html")
        print(f"💡 用浏览器打开 {html_file} 查看对比效果")

def main():
    """主程序"""
    print("=" * 60)
    print("图像预处理效果对比工具")
    print("=" * 60)
    
    # 获取图像路径
    if len(sys.argv) > 1:
        image_path = sys.argv[1]
    else:
        image_path = "Test.png"  # 默认测试图像
    
    if not os.path.exists(image_path):
        print(f"❌ 图像不存在: {image_path}")
        print("用法: python preprocessing_comparison.py <图像路径>")
        return
    
    # 创建处理器并开始处理
    processor = PreprocessingComparison()
    
    # 创建以图像名称命名的输出目录
    base_name = os.path.splitext(os.path.basename(image_path))[0]
    output_dir = f"preprocessing_{base_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    
    results = processor.process_image(image_path, output_dir)
    
    if results:
        print(f"\n🎉 处理完成！")
        print(f"📁 结果目录: {output_dir}")
        print(f"🌐 对比页面: {os.path.join(output_dir, 'comparison.html')}")
        print(f"\n💡 建议:")
        print(f"1. 在浏览器中打开 comparison.html 查看所有效果")
        print(f"2. 观察哪种方法让文字最清晰")
        print(f"3. 选择噪声最少、对比度最好的版本")
        print(f"4. 告诉我哪个效果最好，我会更新代码")

if __name__ == "__main__":
    main()
